package com.zshield.stream;

import com.zshield.proc.OfflineQueryESProc;
import com.zshield.config.KafkaConfig;
import com.zshield.entry.*;
import com.zshield.proc.GenericQueryProc;
import com.zshield.proc.PreProc;
import com.zshield.run.KafkaPreCompute;
import com.zshield.util.KafkaClusterStatus;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.streams.KafkaStreams;
import org.apache.kafka.streams.Topology;
import org.apache.kafka.streams.state.*;

import org.apache.log4j.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public class Stream implements Runnable {
    private Topology topo;
    private StoreBuilder builder1;
    private StoreBuilder builder2;
    private StoreBuilder builder3;
    private StoreBuilder offline;
    private static final Logger logger = Logger.getLogger(Stream.class);

    public Stream() {
        topo = new Topology();
        KeyValueBytesStoreSupplier tsup1 = Stores.inMemoryKeyValueStore("kvDay1");
        KeyValueBytesStoreSupplier tsup2 = Stores.inMemoryKeyValueStore("qstore");
        KeyValueBytesStoreSupplier tsup3 = Stores.inMemoryKeyValueStore("kvDay3");
        KeyValueBytesStoreSupplier tsup4 = Stores.inMemoryKeyValueStore("kvOffLine");
        builder1 = Stores.keyValueStoreBuilder(tsup1, Serdes.String(), Serdes.String());
        builder2 = Stores.keyValueStoreBuilder(tsup2, Serdes.String(), Serdes.String());
        builder3 = Stores.keyValueStoreBuilder(tsup3, Serdes.String(), Serdes.String());
        offline = Stores.keyValueStoreBuilder(tsup4, Serdes.String(), Serdes.String());
    }

    @Override
    public void run() {
        KafkaClusterStatus.isTopicCreated(KafkaPreCompute.input_topic);

        topo.addSource("source", KafkaPreCompute.input_topic)
            .addProcessor("pre_proc", () -> new PreProc(),"source")

            .addProcessor("net_deny_query_proc", () -> new GenericQueryProc<>(NetDenyEntry.class),"pre_proc")
            .addProcessor("netio_query_proc", () -> new GenericQueryProc<>(NetIoEntry.class),"pre_proc")
            .addProcessor("filesystem_inout_query_proc", () -> new GenericQueryProc<>(FileSystemInOutEntry.class),"pre_proc")
            .addProcessor("filesystem_query_proc", () -> new GenericQueryProc<>(FileSystemEntry.class),"pre_proc")
            .addProcessor("filetree_query_proc", () -> new GenericQueryProc<>(FileSystemTreeEntry.class),"pre_proc")
            .addProcessor("content_analysis_query_proc", () -> new GenericQueryProc<>(ContentAnalysisEntry.class),"pre_proc")
            .addProcessor("linux_netio_query_proc", () -> new GenericQueryProc<>(LinuxNetioEntry.class),"pre_proc")

            .addProcessor("offline_net_deny_query_proc", () -> new OfflineQueryESProc(NetDenyEntry.class),"net_deny_query_proc")
            .addProcessor("offline_netio_query_proc", () -> new OfflineQueryESProc(NetIoEntry.class),"netio_query_proc")
            .addProcessor("offline_filesystem_inout_query_proc", () -> new OfflineQueryESProc(FileSystemInOutEntry.class),"filesystem_inout_query_proc")
            .addProcessor("offline_filesystem_query_proc", () -> new OfflineQueryESProc(FileSystemEntry.class),"filesystem_query_proc")
            .addProcessor("offline_filesystem_tree_query_proc", () -> new OfflineQueryESProc(FileSystemTreeEntry.class),"filetree_query_proc")
            .addProcessor("offline_content_analysis_query_proc", () -> new OfflineQueryESProc(ContentAnalysisEntry.class),"content_analysis_query_proc")
            .addProcessor("offline_linuxNetIo_app_query_proc", () -> new OfflineQueryESProc(LinuxNetioEntry.class),"linux_netio_query_proc")

            .addStateStore(builder1,"netio_query_proc","filesystem_query_proc","filesystem_inout_query_proc", "net_deny_query_proc", "filetree_query_proc","content_analysis_query_proc","linux_netio_query_proc")
            .addStateStore(builder2,"netio_query_proc","filesystem_query_proc","filesystem_inout_query_proc", "net_deny_query_proc", "filetree_query_proc","content_analysis_query_proc","linux_netio_query_proc")
            .addStateStore(builder3,"netio_query_proc","filesystem_query_proc","filesystem_inout_query_proc", "net_deny_query_proc", "filetree_query_proc","content_analysis_query_proc","linux_netio_query_proc")
            .addStateStore(offline, "offline_net_deny_query_proc", "offline_netio_query_proc", "offline_filesystem_inout_query_proc", "offline_filesystem_query_proc", "offline_filesystem_tree_query_proc","offline_content_analysis_query_proc","offline_linuxNetIo_app_query_proc")
            .addSink("Generic_query_proc",KafkaPreCompute.medium_topic,"netio_query_proc", "filesystem_query_proc", "filesystem_inout_query_proc", "net_deny_query_proc", "filetree_query_proc","content_analysis_query_proc","linux_netio_query_proc", "offline_net_deny_query_proc", "offline_netio_query_proc", "offline_filesystem_inout_query_proc", "offline_filesystem_query_proc", "offline_filesystem_tree_query_proc","offline_content_analysis_query_proc","offline_linuxNetIo_app_query_proc");

        if (!KafkaPreCompute.is_upgrade) {
            topo.addSink("sensor_output",KafkaPreCompute.output_topic,"source");
        }

        final KafkaStreams streams = new KafkaStreams(topo,KafkaConfig.getStreamConfig());
        final CountDownLatch latch = new CountDownLatch(1);

        streams.setUncaughtExceptionHandler((Thread thread, Throwable e) -> {
            logger.error("[stream caught exception] [The reason for error {" + e + "}");
            KafkaConfig.printErrorLog(logger,e);
            latch.countDown();
        });

        Runtime.getRuntime().addShutdownHook(new Thread(streams::close));

        try {
            streams.start();
            latch.await();
        }catch (Throwable e){
            logger.error("[CountDownLatch await exception] [The reason for error {" + e + "}");
            KafkaConfig.printErrorLog(logger,e);
            System.exit(1);
        }
        streams.close(5,TimeUnit.SECONDS);
        System.exit(0);
    }


}
